from utils import PreprocessorCvcuda ,PostprocessorCvcuda, TRT_Infer
import cvcuda
import cv2
import torch
import numpy as np
import nvcv

class CVCUDA_TensorRT:
    def __init__(self):
        self.prepocess = PreprocessorCvcuda(0)
        self.inference = TRT_Infer("/datas/tky/cvcuda/cvcuda12/yolov8.trt")
        # self.postpocess = PostprocessorCvcuda()
        
    def __call__(self, tensor_in):
        frame_nhwc,resized,rgb_tensor,normalized = self.prepocess(tensor_in,[640,640])
        out_trt = self.inference(normalized)
        print(out_trt.shape,print(type(out_trt)))
        pass
if __name__ == "__main__":
    cvcuda_tensorrt = CVCUDA_TensorRT()
    img = cv2.imread("/datas/tky/cvcuda/cvcuda12/cvcuda-trt/1.png")
    infer_batch = 1
    frame_nhwc_batch = []
    while len(frame_nhwc_batch)<infer_batch:
        frame_nhwc_batch.append(cvcuda.as_tensor(torch.as_tensor(img).to(device=0, non_blocking=True),"HWC"))
    bgr_img_batch = cvcuda.stack(frame_nhwc_batch)
    cvcuda_tensorrt(bgr_img_batch)
    
    

    
    